from QKDNetwork import QKDNetwork
from docplex.mp.model import Model
import time

def getDualPSTimeSlot(routeList, net: QKDNetwork, forOverhead = False):
    # 使用LP求解双端密钥发生器情况下的时隙数量
    net.clearLoad()
    def addFlowOnPath(sd, path): # 施加流
        for i in range(len(path)-1):
            link_start = path[i]
            link_end = path[i + 1]
            net.G[link_start][link_end]["load"] += sd[2]
            
    for i in routeList:
        path = i["path"]
        addFlowOnPath(i["sd"], path)
    
    model = Model(name="dual_ps_time_slot")
    
    p_dim1, p_dim2 = net.num_nodes+1, net.num_nodes+1
    p = model.integer_var_dict(
        [(i, j) for i in range(p_dim1)
                for j in range(p_dim2)], lb=0 ,name='p'
    )
    
    # 最终完成的总时隙
    t = model.integer_var(name="t", lb=0)
    
    for edge in net.G.edges:
        u = edge[0]
        v = edge[1]
        if(net.G.nodes[u]["transmitter"] == 1 and
           net.G.nodes[v]["transmitter"] == 1):
            model.add_constraint(net.G.nodes[u]["transmitter_rate"] * p[u, v] + net.G.nodes[v]["transmitter_rate"] * p[v, u] == net.G[u][v]["load"])
        
    for node in net.G.nodes:
        if(net.G.nodes[node]["transmitter"] == 0):
            continue
        # 该节点是发射器节点
        # 为链路的对端节点不是发射器节点的，计数
        for_single_ps_time_slot = 0
        pu_t = model.linear_expr()
        for neighbor in net.G.neighbors(node):
            if(net.G.nodes[neighbor]["transmitter"] == 0):
                for_single_ps_time_slot += net.G[node][neighbor]["load"] / net.G.nodes[node]["transmitter_rate"]
            else:
                pu_t += p[node, neighbor]
        pu_t += for_single_ps_time_slot
        pu_t -= t
        model.add_constraint(pu_t <= 0)
        
    model.minimize(t)
    # model.export_as_lp("dual_ps.lp")
    start = time.time()
    sol = model.solve()
    end = time.time()
    
    if(forOverhead):
        return end - start
    
    if sol is not None:  # 如果找到可行解
        # print("DualPSTimeSlot", model.get_var_by_name("t").solution_value)
        return model.get_var_by_name("t").solution_value
    else:
        print("DualPSTimeSlot求解失败")
        return 0
        
    